import rospy
from std_msgs.msg import Float64MultiArray
from geometry_msgs.msg import Twist
import cv2
import numpy as np
import yaml

class CoordinateSystemSolver:
    def __init__(self, camera_param_path):
        self.camera_param_path = camera_param_path
        self.camera_matrix, self.dist_coeffs = self.load_camera_params(camera_param_path)
        self.color_ranges = {
            '+x': ([0, 0, 90], [80, 80, 255]),  
            '-x': ([90, 90, 0], [255, 255, 80]),
            '+y': ([0, 90, 90], [80, 255, 255]),
            '-y': ([90, 0, 0], [255, 80, 80]),
            '+z': ([0, 90, 0], [80, 255, 80]),
            '-z': ([90, 0, 90], [255, 80, 255])
        }
        self.ideal_colors = {
            '+x': (0, 0, 255),
            '-x': (255, 255, 0),
            '+y': (0, 255, 255),
            '-y': (255, 0, 0),
            '+z': (0, 255, 0),
            '-z': (255, 0, 255)
        }
        self.initial_coordinates = {
            '+x': [8.5, 0, 0],
            '-x': [-8.5, 0, 0],
            '+y': [0, 8.5, 0],
            '-y': [0, -8.5, 0],
            '+z': [0, 0, 8.5],
            '-z': [0, 0, -8.5]
        }

    def load_camera_params(self, file_path):
        """加载相机标定参数"""
        with open(file_path, "r") as f:
            data = yaml.safe_load(f)
        camera_matrix = np.array(data["camera_matrix"])
        dist_coeffs = np.array(data["distortion_coefficients"])
        return camera_matrix, dist_coeffs

    def points_detect(self, frame):
        """检测图像中的颜色点并返回坐标列表"""
        points = {}
        vis_frame = frame.copy()

        # 使用高斯模糊来平滑图像，减少噪声
        blurred_frame = cv2.GaussianBlur(frame, (11, 11), 0)

        for axis, (lower, upper) in self.color_ranges.items():
            mask = cv2.inRange(blurred_frame, np.array(lower), np.array(upper))
            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            if contours:
                largest_contour = max(contours, key=cv2.contourArea)
                M = cv2.moments(largest_contour)
                if M["m00"] != 0:
                    cX = int(M["m10"] / M["m00"])
                    cY = int(M["m01"] / M["m00"])
                    points[axis] = (cX, cY)
                    cv2.circle(vis_frame, (cX, cY), 5, self.ideal_colors[axis], -1)

        return points, vis_frame

    def point_Coordinate_list_prepare(self, points):
        """准备对应表"""
        coordinate_list = []
        for axis, coord in self.initial_coordinates.items():
            if axis in points:
                coordinate_list.append((coord, points[axis]))
        return coordinate_list

    def point_pnp_solve(self, valid_points):
        """使用有效数据进行 PNP 解算"""
        object_points = []
        image_points = []

        for point in valid_points:
            object_points.append(point[0])
            image_points.append(point[1])

        object_points = np.array(object_points, dtype=np.float32)
        image_points = np.array(image_points, dtype=np.float32)

        success, rvec_cam, tvec_cam = cv2.solvePnP(
            object_points, image_points, self.camera_matrix, self.dist_coeffs,
            flags=cv2.SOLVEPNP_ITERATIVE,
            useExtrinsicGuess=True,
            rvec=np.array([0.0, 0.0, 0.0], dtype=np.float32),
            tvec=np.array([0.0, 0.0, 1.0], dtype=np.float32)
        )

        if success:
            R_cam, _ = cv2.Rodrigues(rvec_cam)
            return R_cam, tvec_cam
        else:
            return None, None

class PnpSolveNode:
    def __init__(self):
        self.solver = CoordinateSystemSolver(rospy.get_param('~camera_param_path'))
        self.detected_point_sub = rospy.Subscriber('detected_point', Float64MultiArray, self.detected_point_callback)
        self.solve_request_pub = rospy.Publisher('solve_request_from_vision', Twist, queue_size=10)

    def detected_point_callback(self, msg):
        points = {'+x': (msg.data[0], msg.data[1]), '+y': (msg.data[2], msg.data[3])}
        valid_points = self.solver.point_Coordinate_list_prepare(points)
        R_cam, tvec_cam = self.solver.point_pnp_solve(valid_points)
        if R_cam is not None and tvec_cam is not None:
            twist_msg = Twist()
            twist_msg.linear.x = tvec_cam[0]
            twist_msg.linear.y = tvec_cam[1]
            twist_msg.linear.z = tvec_cam[2]
            self.solve_request_pub.publish(twist_msg)

if __name__ == '__main__':
    rospy.init_node('pnp_solve_node')
    node = PnpSolveNode()
    rospy.spin()